import os
import pandas as pd
import numpy as np
import time
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.set_option('display.width', 1000)
pd.set_option('display.float_format', lambda x: '%.7f' % x)
np.set_printoptions(suppress=True, precision=4, threshold=np.inf)
report_path = r"C:\Users\huajia\Desktop\rqalpha4\rqalpha\plot_result\my_first_strategy_file9\2022-01-04_2022-01-08"
trades = []
portfolios = []
trade_file = os.path.join(report_path, 'trades.csv')
trade_df = pd.read_csv(trade_file, encoding='gbk')
trade_df['trade_date'] = pd.to_datetime(trade_df['datetime']).dt.strftime('%Y-%m-%d')
# print(trade_df)

group_path = r"C:\Users\huajia\Desktop\rqalpha4\context_df.csv"
group_df = pd.read_csv(group_path, encoding='gbk').reset_index(drop=True).drop(columns=['Unnamed: 0'])
group_sub_df = group_df.iloc[:, 1:]
# print(group_sub_df)
index_li = group_sub_df.index[group_sub_df.isnull().all(1)]
# print(mask)
group_df = group_df[~group_df.index.isin(index_li)]
# print(group_df.head(20))
date_li = group_df['0'].unique().tolist()
date_li = [i for i in date_li if i in trade_df['trade_date'].unique().tolist()]
print(date_li)
# print(group_df)


def cal(n):
    all_df = []
    for i in date_li[0:-1]:
        print('group {}:'.format(n), i)
        sub_df = group_df[group_df['0'] == i]
        # print(sub_df)
        if n >= len(sub_df):
            print('pass----')
            continue
        group_symbols = sub_df.iloc[n, 1:].values.tolist()
        # print(len(group_symbols), group_symbols)

        # group_temp = trade_df.loc[(trade_df['trade_date'] == i) & (trade_df['side'] == 'BUY')]
        # print('group_temp length:', len(group_temp))
        # group_temp = trade_df.loc[(trade_df['trade_date'] == i) & (trade_df['side'] == 'BUY') &
        #                                (trade_df['order_book_id'].isin(group_symbols))]
        # print('group_temp length:', len(group_temp))

        # group_trade_buy = trade_df.loc[(trade_df['trade_date'] == i) &
        #                                (trade_df['order_book_id'].isin(group_symbols)) &
        #                                (trade_df['side'] == 'BUY')]
        group_trade_buy = trade_df.loc[(trade_df['trade_date'] == i) &
                                       (trade_df['side'] == 'BUY')]
        group_trade_buy = group_trade_buy[['order_book_id', 'side', 'last_price', 'trade_date']].reset_index(drop=True)
        # print(group_trade_buy)
        index = date_li.index(i)
        index_next = index + 1
        next_trade_date = date_li[index_next]
        # group_trade_sell = trade_df.loc[(trade_df['trade_date'] == next_trade_date) &
        #                                 (trade_df['order_book_id'].isin(group_symbols)) &
        #                                 (trade_df['side'] == 'SELL')]
        group_trade_sell = trade_df.loc[(trade_df['trade_date'] == next_trade_date) &
                                        (trade_df['side'] == 'SELL')]
        group_trade_sell = group_trade_sell[['order_book_id', 'side', 'last_price', 'trade_date']].\
            reset_index(drop=True)
        # print(group_trade_sell)
        if len(group_trade_buy) != len(group_trade_sell):
            print('shape not equal')
            group_trade_sell = group_trade_sell.drop_duplicates()
            group_trade_buy = group_trade_buy.drop_duplicates()
            print(len(group_trade_buy), len(group_trade_sell))
            # if n == 0 and i == '2022-01-07':
            #     group_trade_sell.to_csv('group_trade_sell.csv')
            #     group_trade_buy.to_csv('group_trade_buy.csv')
        merge_df = pd.merge(group_trade_buy, group_trade_sell, how='inner', on=['order_book_id'])
        merge_df['diff'] = merge_df['last_price_y']/merge_df['last_price_x'] - 1
        avg = np.nanmean(merge_df['diff'])
        print(len(merge_df))
        # print(merge_df)
        # time.sleep(15)
        print([next_trade_date, avg])
        all_df.append([next_trade_date, avg])
        # print([i for i in group_symbols if i not in group_trade['order_book_id'].values.tolist()])
        # time.sleep(500)
    # print(all_df)
    return pd.DataFrame(data=all_df, columns=['the_date', 'avg'])


all_dict = {}
for j in [6]:
    df = cal(j)
    all_dict[j] = df
print(all_dict)

for key, df in all_dict.items():
    df['avg'] = (df['avg'] + 1).cumprod()
    all_dict[key] = df
    df.to_csv('{}.csv'.format(key))
print(all_dict)

fig, ax = plt.subplots()
y_major_locator = MultipleLocator(0.008)
ax.yaxis.set_major_locator(y_major_locator)
plt.xticks(rotation=90, fontsize=7)
for key, item in all_dict.items():
    plt.plot(item['the_date'], item['avg'], label='group {}'.format(key))
plt.legend()
plt.show()




